package main

import (
	"encoding/binary"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"strconv"
	"strings"
	"syscall"
)

var (
	errAddrType      = errors.New("socks addr type not supported")
	errVer           = errors.New("socks version not supported")
	errMethod        = errors.New("socks only support 1 method now")
	errAuthExtraData = errors.New("socks authentication get extra data")
	errReqExtraData  = errors.New("socks request get extra data")
	errCmd           = errors.New("socks command not supported")
)

const (
	socksVer5       = 5
	socksCmdConnect = 1
)

func handShake(conn net.Conn) (err error) {
	const (
		idVer     = 0
		idNmethod = 1
	)
	buf := make([]byte, 258)
	var n int
	if n, err = io.ReadAtLeast(conn, buf, idNmethod+1); err != nil {
		return
	}
	if buf[idVer] != socksVer5 {
		return errVer
	}
	nmethod := int(buf[idNmethod])
	msgLen := nmethod + 2
	if n == msgLen { // handshake done, common case
		// do nothing, jump directly to send confirmation
	} else if n < msgLen { // has more methods to read, rare case
		if _, err = io.ReadFull(conn, buf[n:msgLen]); err != nil {
			return
		}
	} else { // error, should not get extra data
		return errAuthExtraData
	}
	// send confirmation: version 5, no authentication required
	_, err = conn.Write([]byte{socksVer5, 0})
	return
}

//////////////////////////////////////////////////////////////////////
func getRequest(conn net.Conn) (host string, extra []byte, err error) {
	const (
		idVer   = 0
		idCmd   = 1
		idType  = 3
		idIP0   = 4
		idDmLen = 4
		idDm0   = 5

		typeIPv4 = 1
		typeDm   = 3
		typeIPv6 = 4

		lenIPv4   = 3 + 1 + net.IPv4len + 2
		lenIPv6   = 3 + 1 + net.IPv6len + 2
		lenDmBase = 3 + 1 + 1 + 2
	)

	// buf size should at least have the same size with the largest possible
	// request size (when addrType is 3, domain name has at most 256 bytes)
	// 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
	buf := make([]byte, 263)
	var n int
	if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil {
		return
	}
	if buf[idVer] != socksVer5 {
		err = errVer
		return
	}
	if buf[idCmd] != socksCmdConnect {
		err = errCmd
		return
	}

	reqLen := -1
	switch buf[idType] {
	case typeIPv4:
		reqLen = lenIPv4
	case typeIPv6:
		reqLen = lenIPv6
	case typeDm:
		reqLen = int(buf[idDmLen]) + lenDmBase
	default:
		err = fmt.Errorf("addr type %d not supported", buf[idType])
		return
	}

	if n < reqLen { // rare case
		if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil {
			return
		}
	} else if n > reqLen {
		// it's possible to read more than just the request head
		extra = buf[reqLen:n]
	}

	switch buf[idType] {
	case typeIPv4:
		host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
	case typeIPv6:
		host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
	case typeDm:
		host = string(buf[idDm0 : idDm0+buf[idDmLen]])
	}
	// parse port
	port := binary.BigEndian.Uint16(buf[reqLen-2 : reqLen])
	host = net.JoinHostPort(host, strconv.Itoa(int(port)))
	return
}

func handleConnection(conn net.Conn) {
	defer conn.Close()

	closed := false
	defer func() {
		if !closed {
			conn.Close()
		}
	}()

	var err error = nil
	if err = handShake(conn); err != nil {
		fmt.Println("socks handshake:", err)
		return
	}
	host, _, err := getRequest(conn)
	if err != nil {
		fmt.Println("error getting request:", err)
		return
	}
	fmt.Println("remote host:", host)
	_, err = conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x08, 0x43})
	if err != nil {
		fmt.Println("send connection confirmation:", err)
		return
	}
	remote, err := net.Dial("tcp", host)
	if err != nil {
		if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
			log.Println("dial error:", err)
		} else {
			log.Println("error connecting to:", host, err)
		}
		return
	}
	defer func() {
		if !closed {
			remote.Close()
		}
	}()

	//Buffered so that the other goroutine doesn't deadlock
	stopChan := make(chan bool, 1)
	go func() {
		io.Copy(conn, remote)
		stopChan <- true
	}()
	go func() {
		io.Copy(remote, conn)
		stopChan <- true
	}()

	//Wait for either of the copies to stop
	<-stopChan

	return
}

type iplist struct {
	User string `json:"user"`
	Ip   string `json:"ip"`
}

func main() {
	port := flag.Int("port", 1080, "port to listen for connections on")

	flag.Parse()

	address := fmt.Sprintf(":%d", *port)

	ln, err := net.Listen("tcp", address)
	if err != nil {
		log.Printf("Failed to listen on \"%s\" - error: %s", address, err.Error())
		return
	}

	log.Printf("Listening for new connections on %s", address)

	var ips []iplist
	bytes, err := ioutil.ReadFile("./ip.json")
	if err != nil {
		fmt.Println("readFile: ", err.Error())
		return
	}
	err = json.Unmarshal(bytes, &ips)
	if err != nil {
		fmt.Println("Unmarshal: ", err.Error())
		return
	}
	fmt.Println(ips)
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("Failed to accept connection: %s", err.Error())
			continue
		}
		IPandPort := strings.Split(c.RemoteAddr().String(), ":")
		enable := false
		for _, key := range ips {
			if key.Ip == IPandPort[0] {
				enable = true
			}
		}
		if enable {
			go handleConnection(c)
		} else {
			c.Write([]byte("没有权限登陆，请联系管理员"))
			c.Close()
		}
	}
}
